﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Runtime.InteropServices;
using System.IO;
using System.Windows.Forms;

namespace Binalyse.Memory
{
    [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Unicode)]
    public struct PROCESSENTRY32
    {
        public uint Size;
        public uint Usage;
        public uint ProcessID;
        public uint DefaultHeapID;
        public uint ModuleID;
        public uint Threads;
        public uint ParentProcessID;
        public uint PriClassBase;
        public uint Flags;
        [MarshalAs(UnmanagedType.ByValTStr, SizeConst = 260)]
        public string ExeFile;
    }

    [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Unicode)]
    public struct MODULEENTRY32
    {
        public uint Size;
        public uint ModuleID;
        public uint ProcessID;
        public uint GlblcntUsage;
        public uint ProccntUsage;
        public IntPtr BaseAddr;
        public uint modBaseSize;
        public IntPtr hModule;
        [MarshalAs(UnmanagedType.ByValTStr, SizeConst = 256)]
        public string Module;//[MAX_MODULE_NAME32 + 1];
        [MarshalAs(UnmanagedType.ByValTStr, SizeConst = 260)]
        public string ExePath;//[MAX_PATH];
    }

    [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Unicode)]
    public struct MEMORY_BASIC_INFORMATION
    {
        public IntPtr BaseAddress;
        public IntPtr AllocationBase;
        public uint AllocationProtect;
        public uint RegionSize;
        public uint State;
        public uint Protect;
        public uint Type;
    }

    [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Unicode)]
    public struct DEBUG_HEAP_INFORMATION
    {
        public uint Base; /* 0x00 */
        public uint Flags; /*  0x04 */
        public UInt16 Granularity; /*  0x08 */
        public UInt16 Unknown; /*  0x0A */
        public uint Allocated; /*  0x0C */
        public uint Committed; /*  0x10 */
        public uint TagCount; /*  0x14 */
        public uint BlockCount; /*  0x18 */
        [MarshalAs(UnmanagedType.ByValArray, SizeConst = 7)]
        public uint[] Reserved; /*  0x1C */
        public IntPtr Tags; /*  0x38 */
        public IntPtr Blocks; /*  0x3C */
    }

    [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Ansi)]
    public struct DEBUG_MODULE_INFORMATION
    {
        [MarshalAs(UnmanagedType.ByValArray, SizeConst = 2)]
        public uint[] Reserved; /* 0x00 */
        public uint Base; /* 0x08 */
        public uint Size; /* 0x0C */
        public uint Flags; /* 0x10 */
        public UInt16 Index; /* 0x14 */
        public UInt16 Unknown; /* 0x16 */
        public UInt16 LoadCount; /* 0x18 */
        public UInt16 ModuleNameOffset; /* 0x1A */
        [MarshalAs(UnmanagedType.ByValTStr, SizeConst = 256)]
        public string ImageName; /* 0x1C */
    }

    [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Unicode)]
    public struct DEBUG_BUFFER
    {
        public IntPtr SectionHandle;
        public IntPtr SectionBase;
        public IntPtr RemoteSectionBase;
        public uint SectionBaseDelta;
        public IntPtr EventPairHandle;
        [MarshalAs(UnmanagedType.ByValArray, SizeConst = 2)]
        public uint[] Unknown;
        public IntPtr RemoteThreadHandle;
        public uint InfoClassMask;
        public uint SizeOfInfo;
        public uint AllocatedSize;
        public uint SectionSize;
        public IntPtr ModuleInformation;
        public IntPtr BackTraceInformation;
        public IntPtr HeapInformation;
        public IntPtr LockInformation;
        [MarshalAs(UnmanagedType.ByValArray, SizeConst = 8)]
        public IntPtr[] Reserved;
    }
    public class PESection
    {
        public string Name;
        public uint Flag;
        public IntPtr Address;
        public uint Length;
        public uint State;
        public uint Type;
        public uint Protect;
        public byte[] Data;
    }
    public class HeapBlock
    {
        public int StartIndex;
        public int Length;
        public uint Flag;
    }
    public class Heap
    {
        public IntPtr Base;
        public uint Flag;
        public uint Length;
        public uint State;
        public uint Type;
        public uint Protect;
        public byte[] Data;
        public List<HeapBlock> Blocks;
    }
    public static class Win32API
    {
        public static IntPtr INVALID_HANDLE_VALUE = new IntPtr(-1);

        [DllImport("Kernel32.dll", EntryPoint = "CloseHandle")]
        public static extern bool CloseHandle(IntPtr handle);
        [DllImport("Kernel32.dll", EntryPoint = "GetLastError")]
        public static extern uint GetLastError();
        [DllImport("Kernel32.dll", EntryPoint = "GetProcessId")]
        public static extern uint GetProcessId(IntPtr hProcess);

        #region ToolHelp
        
        #region CreateToolhelp32Snapshot
        [DllImport("Kernel32.dll", SetLastError = true)]
        public static extern IntPtr CreateToolhelp32Snapshot(uint dwFlags, uint th32ProcessID);
        public const uint TH32CS_INHERIT = 0x80000000;
        public const uint TH32CS_SNAPHEAPLIST = 0x00000001;
        public const uint TH32CS_SNAPMODULE = 0x00000008;
        public const uint TH32CS_SNAPMODULE32 = 0x00000010;
        public const uint TH32CS_SNAPPROCESS = 0x00000002;
        public const uint TH32CS_SNAPTHREAD = 0x00000004;
        public const uint TH32CS_SNAPALL = TH32CS_INHERIT | TH32CS_SNAPHEAPLIST | TH32CS_SNAPMODULE | TH32CS_SNAPMODULE32 | TH32CS_SNAPPROCESS | TH32CS_SNAPTHREAD;
        #endregion

        [DllImport("Kernel32.dll", EntryPoint = "Process32FirstW")]
        public static extern bool Process32First(IntPtr hSnapshot, ref PROCESSENTRY32 lppe);

        [DllImport("Kernel32.dll", EntryPoint = "Process32NextW")]
        public static extern bool Process32Next(IntPtr hSnapshot, ref PROCESSENTRY32 lppe);

        [DllImport("Kernel32.dll", EntryPoint = "Module32FirstW")]
        public static extern bool Module32First(IntPtr hSnapshot, ref MODULEENTRY32 lppe);

        [DllImport("Kernel32.dll", EntryPoint = "Module32NextW")]
        public static extern bool Module32Next(IntPtr hSnapshot, ref MODULEENTRY32 lppe);

        #endregion

        #region OpenProcess
        [DllImport("Kernel32.dll", EntryPoint = "OpenProcess")]
        public static extern IntPtr OpenProcess(uint dwDesiredAccess, bool bInheritHandle, uint dwProcessId);
        //Required to create a process. 
        public const uint PROCESS_CREATE_PROCESS = (0x0080);
        //Required to create a thread. 
        public const uint PROCESS_CREATE_THREAD = (0x0002);
        //Required to duplicate a handle using DuplicateHandle.  
        public const uint PROCESS_DUP_HANDLE = (0x0040);
        //Required to retrieve certain information about a process, such as its token, exit code, 
        //and priority class (see OpenProcessToken, GetExitCodeProcess, GetPriorityClass, and IsProcessInJob).  
        public const uint PROCESS_QUERY_INFORMATION = (0x0400);
        //Required to retrieve certain information about a process (see QueryFullProcessImageName). 
        //A handle that has the PROCESS_QUERY_INFORMATION access right is automatically granted PROCESS_QUERY_LIMITED_INFORMATION. 
        //Windows Server 2003 and Windows XP/2000:  This access right is not supported.  
        public const uint PROCESS_QUERY_LIMITED_INFORMATION = (0x1000);
        //Required to set certain information about a process, such as its priority class (see SetPriorityClass).  
        public const uint PROCESS_SET_INFORMATION = (0x0200);
        //Required to set memory limits using SetProcessWorkingSetSize.
        public const uint PROCESS_SET_QUOTA = (0x0100);
        //Required to suspend or resume a process. 
        public const uint PROCESS_SUSPEND_RESUME = (0x0800);
        //Required to terminate a process using TerminateProcess.  
        public const uint PROCESS_TERMINATE = (0x0001);
        //Required to perform an operation on the address space of a process (see VirtualProtectEx and WriteProcessMemory).  
        public const uint PROCESS_VM_OPERATION = (0x0008);
        //Required to read memory in a process using ReadProcessMemory.  
        public const uint PROCESS_VM_READ = (0x0010);
        //Required to write to memory in a process using WriteProcessMemory.  
        public const uint PROCESS_VM_WRITE = (0x0020);
        // Required to wait for the process to terminate using the wait functions.  
        public const uint SYNCHRONIZE = (0x00100000);
        #endregion

        #region Read and Write process
        
        [DllImport("Kernel32.dll", EntryPoint = "ReadProcessMemory")]
        public static extern bool ReadProcessMemory(IntPtr hProcess, IntPtr lpBaseAddress, [In,Out] byte[] lpBuffer, uint nSize, ref uint lpNumberOfBytesRead);

        [DllImport("Kernel32.dll", EntryPoint = "WriteProcessMemory")]
        public static extern bool WriteProcessMemory(IntPtr hProcess, IntPtr lpBaseAddress, [In] byte[] lpBuffer, uint nSize, ref uint lpNumberOfBytesWritten);
        
        #endregion

        #region VirtualQueryEx
        [DllImport("Kernel32.dll", EntryPoint = "VirtualQueryEx")]
        public static extern uint VirtualQueryEx(IntPtr hProcess, IntPtr lpAddress, ref MEMORY_BASIC_INFORMATION lpBuffer, uint dwLength);

        #region State
        // Indicates committed pages for which physical storage has been allocated, either in memory or in the paging file on disk.
        public const uint MEM_COMMIT = 0x1000;
        //Indicates free pages not accessible to the calling process and available to be allocated. 
        //For free pages, the information in the AllocationBase, AllocationProtect, Protect, and Type members is undefined. 
        public const uint MEM_FREE = 0x10000;
        //Indicates reserved pages where a range of the process's virtual address space is reserved without any physical storage being allocated. 
        //For reserved pages, the information in the Protect member is undefined. 
        public const uint MEM_RESERVE = 0x2000;
        #endregion
        
        #region Protect
        //以下选项单独使用
        public const uint PAGE_EXECUTE = 0x10;
        public const uint PAGE_EXECUTE_READ = 0x20;
        public const uint PAGE_EXECUTE_READWRITE = 0x40;
        public const uint PAGE_EXECUTE_WRITECOPY = 0x80;
        public const uint PAGE_NOACCESS = 0x01;
        public const uint PAGE_READONLY = 0x02;
        public const uint PAGE_READWRITE = 0x04;
        public const uint PAGE_WRITECOPY = 0x08;
        //以下选项可与上面的配合使用
        public const uint PAGE_GUARD = 0x100;
        public const uint PAGE_NOCACHE = 0x200;
        public const uint PAGE_WRITECOMBINE = 0x400;
        #endregion
        
        #region Type
        // Indicates that the memory pages within the region are mapped into the view of an image section.
        public const uint MEM_IMAGE = 0x1000000;
        // Indicates that the memory pages within the region are mapped into the view of a section.
        public const uint MEM_MAPPED = 0x40000;
        //Indicates that the memory pages within the region are private (that is, not shared by other processes).
        public const uint MEM_PRIVATE = 0x20000;
        #endregion

        #endregion

        #region RTL functions

        public const uint PDI_MODULES = 0x00000001;
        public const uint PDI_BACKTRACE = 0x00000002;
        public const uint PDI_HEAPS = 0x00000004;
        public const uint PDI_HEAPS_TAGS = 0x00000008;
        public const uint PDI_HEAP_BLOCKS = 0x00000010;
        public const uint PDI_LOCKS = 0x00000020;

        //RtlCreateQueryDebugBuffer 返回类型是DEBUG_BUFFER，这里用IntPtr，如果要转换为DEBUG_BUFFER，可使用Marshal.PtrToStructure方法
        [DllImport("NtDLL.dll", EntryPoint = "RtlCreateQueryDebugBuffer")]
        public static extern IntPtr RtlCreateQueryDebugBuffer(int nZero, bool bFalse);

        [DllImport("NtDLL.dll", EntryPoint = "RtlQueryProcessDebugInformation")]
        public static extern uint RtlQueryProcessDebugInformation(uint ProcessId, uint pdiOption, IntPtr pDebugBuffer);

        [DllImport("NtDLL.dll", EntryPoint = "RtlDestroyQueryDebugBuffer")]
        public static extern uint RtlDestroyQueryDebugBuffer(IntPtr pDebugBuffer);

        #endregion

        #region GetPESections
        public const uint PE_SIGN = 0x00004550;
        public const int PE_HEADER_LEN = 0x18;
        public const int PE_SECTION_MAX_COUNT = 0x60;
        public const int IMAGE_SIZEOF_SECTION_HEADER = 0x28;
        public static List<PESection> GetPESections(IntPtr hProcess, IntPtr dwBaseAddress,uint readFilter)
        {
            List<PESection> ret = null;
            const int nPageSize = 4096;//4kb页面文件大小
            uint nRead = 0;
            int nPeHeaderOffset, nStartIndexOfPeSec;
            short nNumOfSections, nSizeOfOptionalHeader;
            byte[] buffer = new byte[nPageSize];
            if (ReadProcessMemory(hProcess, dwBaseAddress,buffer, nPageSize,ref nRead))
            {
                //验证是否为DOS头
                if (buffer[0] == 0x4D && buffer[1] == 0x5A)//"MZ"
                {
                    nPeHeaderOffset = BitConverter.ToInt32(buffer, 0x3c);
                    if (nPeHeaderOffset < nRead)
                    {
                        //验证PE头
                        if (PE_SIGN == BitConverter.ToUInt32(buffer, nPeHeaderOffset))
                        {
                            nNumOfSections = BitConverter.ToInt16(buffer, nPeHeaderOffset + 6);
                            nSizeOfOptionalHeader = BitConverter.ToInt16(buffer, nPeHeaderOffset + 20);
                            nStartIndexOfPeSec = nPeHeaderOffset + PE_HEADER_LEN + nSizeOfOptionalHeader;
                            //确保访问不越界
                            if(nNumOfSections <= PE_SECTION_MAX_COUNT
                                && nStartIndexOfPeSec > 0
                                && nStartIndexOfPeSec < nRead - nNumOfSections * IMAGE_SIZEOF_SECTION_HEADER)
                            {
                                //typedef struct _IMAGE_SECTION_HEADER {
                                //    BYTE    Name[IMAGE_SIZEOF_SHORT_NAME]; // 00
                                //    union {
                                //            DWORD   PhysicalAddress;
                                //            DWORD   VirtualSize;
                                //    } Misc;                                // 08
                                //    DWORD   VirtualAddress;                // 0C
                                //    DWORD   SizeOfRawData;                 // 10 
                                //    DWORD   PointerToRawData;              // 14
                                //    DWORD   PointerToRelocations;          // 18
                                //    DWORD   PointerToLinenumbers;          // 1C
                                //    WORD    NumberOfRelocations;           // 20
                                //    WORD    NumberOfLinenumbers;           // 22
                                //    DWORD   Characteristics;               // 24
                                //} IMAGE_SECTION_HEADER, *PIMAGE_SECTION_HEADER;
                                MEMORY_BASIC_INFORMATION mbi = new MEMORY_BASIC_INFORMATION();
                                ret = new List<PESection>();
                                for (int index = 0; index < nNumOfSections; index++)
                                {
                                    PESection pe = new PESection();
                                    //进行UTF8解码时会加入大量的\0结束符，而string.Length属性居然将这些结束符全部计算在内，导致字符串合并时存在问题!
                                    string tmpStr = Encoding.UTF8.GetString(buffer, nStartIndexOfPeSec + index * IMAGE_SIZEOF_SECTION_HEADER, 8);
                                    //pe.Name = pe.Name.Substring(0, pe.Name.Length);
                                    char[] chars = tmpStr.ToCharArray();
                                    StringBuilder sb = new StringBuilder(8);
                                    for (int i = 0; i < 8; i++)
                                    {
                                        sb.Append(chars[i] == '\0' ? ' ' : chars[i]);
                                    }
                                    pe.Name = sb.ToString();
                                    pe.Length = BitConverter.ToUInt32(buffer, nStartIndexOfPeSec + index * IMAGE_SIZEOF_SECTION_HEADER + 0x10);
                                    pe.Address = dwBaseAddress +
                                        BitConverter.ToInt32(buffer, nStartIndexOfPeSec + index * IMAGE_SIZEOF_SECTION_HEADER + 0x0C);
                                    pe.Flag = BitConverter.ToUInt32(buffer, nStartIndexOfPeSec + index * IMAGE_SIZEOF_SECTION_HEADER + 0x24);
                                    if (VirtualQueryEx(hProcess, pe.Address, ref mbi, (uint)Marshal.SizeOf(mbi)) != 0)
                                    {
                                        pe.Protect = mbi.Protect;
                                        pe.State = mbi.State;
                                        pe.Type = mbi.Type;
                                        if ((pe.Protect & readFilter) != 0)
                                        {
                                            pe.Data = new byte[mbi.RegionSize];
                                            ReadProcessMemory(hProcess, pe.Address, pe.Data, mbi.RegionSize, ref nRead);
                                            pe.Length = nRead;
                                            ret.Add(pe);
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }
            return ret;
        }
        #endregion

        #region GetHeapInformation
        public static List<Heap> GetHeapInformation(IntPtr hProcess)
        {
            List<Heap> ret = null;
            uint pid = GetProcessId(hProcess);
            uint nRead = 0;
            IntPtr pDebugBuffer = RtlCreateQueryDebugBuffer(0, false);
            if (pDebugBuffer != null)
            {
                if (RtlQueryProcessDebugInformation(pid, PDI_HEAPS, pDebugBuffer) >= 0)
                {
                    DEBUG_HEAP_INFORMATION dhi = new DEBUG_HEAP_INFORMATION();
                    MEMORY_BASIC_INFORMATION mbi = new MEMORY_BASIC_INFORMATION();
                    int nCount = 0;
                    ret = new List<Heap>();

                    DEBUG_BUFFER db = (DEBUG_BUFFER)Marshal.PtrToStructure(pDebugBuffer, typeof(DEBUG_BUFFER));
                    nCount = Marshal.ReadInt32(db.HeapInformation, 0);
                    for (int i = 0; i < nCount; i++)
                    {
                        dhi = (DEBUG_HEAP_INFORMATION)Marshal.PtrToStructure(
                            db.HeapInformation + i * Marshal.SizeOf(dhi) + 4, typeof(DEBUG_HEAP_INFORMATION));
                        IntPtr baseAddress = new IntPtr(dhi.Base);
                        if (VirtualQueryEx(hProcess, baseAddress, ref mbi, (uint)Marshal.SizeOf(mbi)) != 0)
                        {
                            Heap heap = new Heap();
                            heap.Base = baseAddress;
                            heap.Flag = dhi.Flags;
                            heap.State = mbi.State;
                            heap.Type = mbi.Type;
                            //有一部分的堆内存属性是PAGE_EXECUTE_READWRITE，怎么回事？
                            heap.Protect = mbi.Protect;
                            heap.Data = new byte[mbi.RegionSize];
                            ReadProcessMemory(hProcess, heap.Base, heap.Data, mbi.RegionSize, ref nRead);
                            heap.Length = mbi.RegionSize;
                            ret.Add(heap);
                        }
                    }
                }
                RtlDestroyQueryDebugBuffer(pDebugBuffer);
            }
            return ret;
        }
        #endregion

        public static bool AnalyseHeapBlocks(Heap heap)
        {
            heap.Blocks = new List<HeapBlock>();
            for (uint nStartOffset = 0x680; nStartOffset < heap.Length; )
            {
                HeapBlock block = new HeapBlock();
                block.StartIndex = (int)nStartOffset + 8;

                block.Flag = heap.Data[nStartOffset + 5];
                //Block的大小包含了前面的8byte的头在内,实际数据长度应减去一个头的长度
                block.Length = (BitConverter.ToUInt16(heap.Data, (int)nStartOffset) - 1) * 8;
                if (block.Length <= 0)
                {
                    MessageBox.Show(
                        string.Format("分析堆内存{0:X8}时出错!\r\n偏移位于:{1:X8},头数据{2}",
                                        heap.Base.ToInt32(), nStartOffset,
                                        BitConverter.ToString(heap.Data, (int)nStartOffset, 8)),
                                        "错误", MessageBoxButtons.OK, MessageBoxIcon.Error);
                    return false;
                }

                heap.Blocks.Add(block);
                nStartOffset += (uint)block.Length + 8;
            }
            return true;
        }
    }
}
